from __future__ import division
import cv2
import ip_tools as ip
import os
import csv

class DataProcessor:
    """Prepares/loads data as a set of ([features], class])."""

    def __init__(self, dataPath='./', alphabets=[], dumpFolder='', fileExtension='.png'):
        self.__dataPath = self.__fixPath(dataPath)
        self.__alphabets = alphabets
        self.__dumpPath = self.__dataPath + dumpFolder + '/'
        if not os.path.exists(self.__dumpPath):
            os.makedirs(self.__dumpPath)
        self.__dumpFile = self.__dumpPath + 'extracted_data.csv'
        self.__fileExtension = fileExtension
        
    def __fixPath(self, path):
        """If a path does not ends with '/', concatenates '/' to its end, and returns the fixed path."""

        if path[-1] != '/':
            path = path + '/'
        return path

    def __getFeaturesFromImage(self, img):
        """Extracts useful features from the image and returns it.

        x1 = sum of values of all pixels
        x2 = sum of values of all pixels that lie to the left of the central vertical line
        x3 = sum of values of all pixels that lie to the right of the central vertical line
        x4 = sum of values of all pixels that lie above the central horizontal line
        x5 = sum of values of all pixels that lie below the central horizontal line
        x6 = ratio of height to width

        """

        img = ip.resizeTo100x100(img)
        img_zt = ip.convertToZeroTen(img)
        height, width = img_zt.shape
        x1 = 0
        x2 = 0
        x3 = 0
        x4 = 0
        x5 = 0
        topLimit = -1
        bottomLimit = -1
        leftLimit = width + 1
        rightLimit = -1
        for i in range(height):
            for j in range(width):
                if img_zt[i, j] == 10:
                    if topLimit == -1:
                        topLimit = i
                    if bottomLimit < i:
                        bottomLimit = i
                    if rightLimit < j:
                        rightLimit = j
                    if leftLimit > j:
                        leftLimit = j
                x1 = x1 + img_zt[i, j]
                if j < (width / 2):
                    x2 = x2 + img_zt[i, j]
                else:
                    x3 = x3 + img_zt[i, j]
                if i < (height / 2):
                    x4 = x4 + img_zt[i, j]
                else:
                    x5 = x5 + img_zt[i, j]
        x6 = (bottomLimit - topLimit) / (rightLimit - leftLimit)
        return [x1, x2, x3, x4, x5, x6]

    def __dumpExtractedData(self, contents):
        """Dumps the extracted data to 'extracted_data.csv' in the specified dump-path."""

        outfile = open(self.__dumpFile, 'w')
        with outfile:
            outfile.write('x1,x2,x3,x4,x5,x6,alphabet\n')
            for row in contents:
                row_str = ','.join([str(x) for x in row]) + '\n'
                outfile.write(row_str)

    def extractFeaturesFromImageFile(self, imageFile):
        """Reads a greyscale image, and extracts and returns useful features from it."""

        img = ip.readGreyscaleImage(imageFile)
        return self.__getFeaturesFromImage(img)

    def extractFeaturesListFromFolder(self, folderPath):
        """Compiles the list of feature-vectors for all files in the given path with the given extension."""

        featuresList = []
        path = self.__fixPath(folderPath)
        files = [x for x in os.listdir(path) if x.endswith(self.__fileExtension)]
        for imageFile in files:
            featuresList.append(self.extractFeaturesFromImageFile(path + imageFile))
        return featuresList

    def extractData(self):
        """Reads all images for various alphabet-folders, creates a single featuresList annotated with
        alphabet-label, dumps it on a file at the specified path, and returns it."""

        completeList = []
        for alphabet in self.__alphabets:
            folderPath = self.__dataPath + alphabet + "/"
            if not os.path.exists(folderPath):
                continue
            featuresList = self.extractFeaturesListFromFolder(folderPath)
            for row in featuresList:
                row.append(alphabet)  # annotating each row with alphabet-label
            completeList = completeList + featuresList
        self.__dumpExtractedData(completeList)
        return completeList

    def loadData(self):
        """Reads 'extracted_data.csv' in dump-folder under data-path and returns the loaded data."""

        infile = open(self.__dumpFile)
        contents = []
        with infile:
            c = csv.reader(infile, delimiter=',')
            first = True
            for row in c:
                if first:
                    first = False
                else:
                    rowData = [float(x) for x in row[:-1]] + [row[-1]]  # the last field is the alphabet-label
                    contents.append(rowData)
        return contents